﻿#include <learn/math/graph.h>

using namespace xi;
using namespace xi::stl;
using namespace xi::math;
using namespace xi::math::details;

Graph_Sparse::Graph_Sparse(std::size_t n)
{
    for (std::size_t i = 0; i < n; ++i)
    {
        m_adjList.push_back(Link_List{new List<std::size_t>{}});
        m_vData.push_back(VData{});
    }
}

std::size_t Graph_Sparse::size() const
{
    return m_adjList.size();
}

void Graph_Sparse::create_undirected_edge(std::size_t u, std::size_t v, float weight)
{
    m_adjList[u]->push_front(v);
    m_adjList[v]->push_front(u);
    m_eData.insert(std::make_pair(std::make_pair(u, v), EData{weight}));
    m_eData.insert(std::make_pair(std::make_pair(v, u), EData{weight}));
}

void Graph_Sparse::create_directed_edge(std::size_t u, std::size_t v, float weight)
{
    m_adjList[u]->push_front(v);
    m_eData.insert(std::make_pair(std::make_pair(u, v), EData{weight}));
}

void Graph_Sparse::create_undirected_edges(std::size_t u, const std::vector<std::size_t> &v,
                                           const std::vector<float> &weights)
{
    for (std::size_t i = 0; i < v.size(); ++i)
        create_undirected_edge(u, v[i], weights[i]);
}

void Graph_Sparse::create_directed_edges(std::size_t u, const std::vector<std::size_t> &v,
                                         const std::vector<float> &weights)
{
    for (std::size_t i = 0; i < v.size(); ++i)
        create_directed_edge(u, v[i], weights[i]);
}

void Graph_Sparse::transpose()
{
    std::vector<Link_List> new_adjList(size());
    for (std::size_t i = 0; i < size(); ++i)
        new_adjList[i] = Link_List{new List<std::size_t>{}};

    for (std::size_t i = 0; i < size(); ++i)
    {
        auto adj = m_adjList[i];
        for (auto jt = adj->begin(); jt != adj->end(); ++jt)
            new_adjList[*jt]->push_front(i);
    }

    decltype(m_eData) new_eData;
    for (auto &&[pair, data] : m_eData)
    {
        auto [u, v] = pair;
        new_eData.insert(std::make_pair(std::make_tuple(v, u), data));
    }

    m_adjList = std::move(new_adjList);
    m_eData = std::move(new_eData);
}

void Graph_Sparse::remove_duplicates()
{
    for (std::size_t i = 0; i < size(); ++i)
    {
        m_adjList[i]->remove_duplicates();
        m_adjList[i]->remove_if([i](std::size_t v) { return i == v; });
    }
}

std::vector<std::size_t> Graph_Sparse::adjacent_vertices(std::size_t u) const
{
    std::vector<std::size_t> adj;
    for (auto it = m_adjList.at(u)->begin(); it != m_adjList.at(u)->end(); ++it)
        adj.push_back(*it);
    return adj;
}

bool Graph_Sparse::adjacent_edge(std::size_t u, std::size_t v) const
{
    return m_adjList.at(u)->find(v) != m_adjList.at(u)->end();
}

std::unique_ptr<MN_Tree<float *>> Graph_Sparse::bfs(std::size_t s)
{
    // 初始化
    for (auto &&data : m_vData)
    {
        data.color = VData::Color::White;
        data.data = std::numeric_limits<float>::infinity();
    }

    m_vData.at(s).color = VData::Color::Gray;
    m_vData.at(s).data = 0;

    std::queue<std::size_t> q;
    q.push(s);

    // 构建搜索树
    auto tree = std::make_unique<MN_Tree<float *>>();
    std::queue<MN_Tree<float *>::Node *> q_tree;
    q_tree.push(tree->insert(&m_vData.at(s).data));

    while (!q.empty())
    {
        auto u = q.front();
        q.pop();

        auto node = q_tree.front();
        q_tree.pop();

        for (auto it = m_adjList.at(u)->begin(); it != m_adjList.at(u)->end(); ++it)
        {
            auto v = *it;
            if (m_vData.at(v).color == VData::Color::White)
            {
                m_vData.at(v).color = VData::Color::Gray;
                m_vData.at(v).data = m_vData.at(u).data + 1;
                q.push(v);
                q_tree.push(tree->insert(&m_vData.at(v).data, node));
            }
        }

        m_vData.at(u).color = VData::Color::Black;
    }

    return tree;
}

std::vector<std::tuple<float, float>> Graph_Sparse::dfs(std::size_t s)
{
    std::vector<std::tuple<float, float>> intervals(size());

    // 初始化
    for (auto &&data : m_vData)
    {
        data.color = VData::Color::White;
        data.data = 0;
    }

    std::size_t count = 0;
    std::function<void(std::size_t u)> dfs_visit = [&](std::size_t u) {
        count++;
        intervals.at(u) = {count, 0};
        m_vData.at(u).color = VData::Color::Gray;
        for (auto it = m_adjList.at(u)->begin(); it != m_adjList.at(u)->end(); ++it)
            if (m_vData.at(*it).color == VData::Color::White)
                dfs_visit(*it);
        m_vData.at(u).color = VData::Color::Black;
        count++;
        intervals.at(u) = {std::get<0>(intervals.at(u)), count};
    };

    dfs_visit(s);
    return intervals;
}

std::ostream &xi::math::operator<<(std::ostream &os, const Graph_Sparse &graph)
{
    for (std::size_t i = 0; i < graph.m_adjList.size(); ++i)
    {
        os << "(" << i << ", " << graph.m_vData.at(i) << ")";

        auto adj = graph.m_adjList.at(i);
        for (auto jt = adj->begin(); jt != adj->end(); ++jt)
        {
            os << " -> "
               << "(" << *jt << ", " << graph.m_eData.at(std::make_pair(i, *jt)) << ")";
        }
        os << std::endl;
    }
    return os;
}

Graph_Dense::Graph_Dense(std::size_t n) : m_adjMat(n, n), m_wMat(n, n)
{
    m_adjMat.set_zero();
    m_wMat.set_zero();
}

std::size_t Graph_Dense::size() const
{
    return m_adjMat.rows();
}

void Graph_Dense::create_undirected_edge(std::size_t u, std::size_t v, float weight)
{
    m_adjMat(u, v) += 1;
    m_adjMat(v, u) += 1;
    m_wMat(u, v) = weight;
    m_wMat(v, u) = weight;
}

void Graph_Dense::create_directed_edge(std::size_t u, std::size_t v, float weight)
{
    m_adjMat(u, v) += 1;
    m_wMat(u, v) = weight;
}

void Graph_Dense::create_undirected_edges(std::size_t u, const std::vector<std::size_t> &v,
                                          const std::vector<float> &weights)
{
    for (std::size_t i = 0; i < v.size(); ++i)
        create_undirected_edge(u, v[i], weights[i]);
}

void Graph_Dense::create_directed_edges(std::size_t u, const std::vector<std::size_t> &v,
                                        const std::vector<float> &weights)
{
    for (std::size_t i = 0; i < v.size(); ++i)
        create_directed_edge(u, v[i], weights[i]);
}

void Graph_Dense::transpose()
{
    m_adjMat.transpose();
    m_wMat.transpose();
}

void Graph_Dense::remove_duplicates()
{
    for (std::size_t i = 0; i < m_adjMat.rows(); ++i)
    {
        m_adjMat(i, i) = 0;
        m_wMat(i, i) = 0;
        for (std::size_t j = 0; j < m_adjMat.cols(); ++j)
        {
            if (m_adjMat(i, j) > 1)
                m_adjMat(i, j) = 1;
        }
    }
}

std::vector<std::size_t> Graph_Dense::adjacent_vertices(std::size_t u) const
{
    std::vector<std::size_t> adj;
    for (std::size_t i = 0; i < m_adjMat.rows(); ++i)
    {
        if (m_adjMat(u, i) > 0)
            adj.push_back(i);
    }
    return adj;
}

bool Graph_Dense::adjacent_edge(std::size_t u, std::size_t v) const
{
    return m_adjMat(u, v) > 0;
}

std::size_t Graph_Dense::shortest_path(std::size_t u, std::size_t v, std::size_t len)
{
    if (u == v)
        return 0;

    auto mat = m_adjMat;
    std::size_t count = 1;
    while (mat(u, v) == 0 && count < len)
    {
        ++count;
        mat = mat * m_adjMat;
    }
    return mat(u, v) > 0 ? count : -1;
}

std::ostream &xi::math::operator<<(std::ostream &os, const Graph_Dense &graph)
{
    for (std::size_t i = 0; i < graph.m_adjMat.rows(); ++i)
    {
        os << i;
        for (std::size_t j = 0; j < graph.m_adjMat.cols(); ++j)
        {
            if (graph.m_adjMat(i, j) > 0)
                os << " -> "
                   << "(" << j << " x " << graph.m_adjMat(i, j) << ", " << graph.m_wMat(i, j) << "f) ";
        }
        os << std::endl;
    }
    return os;
}